# coding:utf-8
"""
断裂点空间分布分析脚本
功能:
1. 读取指定任务目录下graph_series/最大timestep的网络,提取末态断裂点
2. 绘制yz, xy, xz平面热力图,以及x, y, z方向分布直方图
3. 所有图片和同名txt数据文件保存在breakpoint_distribution_output子文件夹,txt文件分为append data和main data两部分
4. 增加调试输出,便于排查断裂点提取问题
"""
import os
import pickle
import numpy as np
import matplotlib.pyplot as plt
import networkx as nx
import GraphTools.basic_method as bm
import MDAnalysis as mda
import GraphTools.graph_method as gm

# 自动检测节点坐标字段


def auto_detect_pos_key(graph):
    for key in ['posi', 'position', 'pos']:
        if key in next(iter(graph.nodes(data=True)))[1]:
            return key
    raise ValueError('No valid position key found in node attributes.')

# 获取graph_series下最大timestep的graph.pkl路径


def get_last_graph_path(graph_series_dir):
    files = [f for f in os.listdir(
        graph_series_dir) if f.endswith('-graph.pkl')]
    if not files:
        raise FileNotFoundError('No *-graph.pkl found in graph_series.')
    # 提取timestep
    max_t = max(int(f.split('-')[0]) for f in files)
    return os.path.join(graph_series_dir, f'{max_t}-graph.pkl')

# 修正: 支持MultiGraph多重边遍历,正确提取断裂边


def extract_breakpoints_by_compare(init_graph, last_graph, pos_key, debug=False):
    """
    通过对比初始帧和末态帧的边集合，找出断裂边，并返回断裂点中点坐标。
    """
    # 支持MultiGraph，需用(u,v,k)唯一标识
    init_edges = set(init_graph.edges(keys=True))
    last_edges = set(last_graph.edges(keys=True))
    broken_edges = init_edges - last_edges
    mids = []
    for u, v, k in broken_edges:
        p0 = init_graph.nodes[u][pos_key]
        p1 = init_graph.nodes[v][pos_key]
        box = bm.get_box_dimensions(init_graph.graph['box'])
        mids.append(tuple(pbc_midpoint(p0, p1, box)))
    if debug:
        print(f"初始边数: {len(init_edges)}")
        print(f"末态边数: {len(last_edges)}")
        print(f"断裂边数: {len(broken_edges)}")
        for i, (u, v, k) in enumerate(list(broken_edges)[:5]):
            print(
                f"断裂边 {i}: {u}-{v}-{k}, pos0={init_graph.nodes[u][pos_key]}, pos1={init_graph.nodes[v][pos_key]}")
        if len(broken_edges) == 0:
            print("Warning: No broken edges found. All output will be empty.")
    return np.array(mids)  # shape (N,3)

# 保存数据到txt,严格按规范


def save_data_txt(filename, append_dict, main_data, main_header, empty_note=None):
    with open(filename, 'w') as f:
        f.write('# append data\n')
        for k, v in append_dict.items():
            f.write(f'{k} = {v}\n')
        f.write('# main data\n')
        f.write(','.join(main_header) + '\n')
        if empty_note:
            f.write(f'# {empty_note}\n')
        for row in main_data:
            f.write(','.join(map(str, row)) + '\n')

# 绘制热力图


def plot_heatmap(mids, x_idx, y_idx, x_label, y_label, out_prefix, out_dir, bins=60):
    out_png = os.path.join(out_dir, f'{out_prefix}_heatmap.png')
    out_txt = os.path.join(out_dir, f'{out_prefix}_heatmap.txt')
    if len(mids) == 0:
        H = np.zeros((bins, bins))
        xedges = yedges = np.linspace(0, 1, bins+1)
        empty_note = 'No breakpoints found.'
    else:
        x = mids[:, x_idx]
        y = mids[:, y_idx]
        xedges = np.linspace(x.min(), x.max(), bins+1)
        yedges = np.linspace(y.min(), y.max(), bins+1)
        H, _, _ = np.histogram2d(x, y, bins=[xedges, yedges])
        empty_note = None
    # 画图
    plt.figure(figsize=(6, 5))
    plt.imshow(H.T, origin='lower', aspect='auto',
               extent=[xedges[0], xedges[-1], yedges[0], yedges[-1]],
               cmap='viridis')
    plt.xlabel(x_label)
    plt.ylabel(y_label)
    plt.title(f'Breakpoints {x_label}{y_label} Heatmap')
    plt.colorbar(label='Count')
    plt.tight_layout()
    plt.savefig(out_png, dpi=300)
    plt.close()
    # 保存数据
    # append data: mean, std
    append_dict = {
        f'mean_{x_label.lower()}': np.mean(mids[:, x_idx]) if len(mids) else 0,
        f'std_{x_label.lower()}': np.std(mids[:, x_idx]) if len(mids) else 0,
        f'mean_{y_label.lower()}': np.mean(mids[:, y_idx]) if len(mids) else 0,
        f'std_{y_label.lower()}': np.std(mids[:, y_idx]) if len(mids) else 0,
    }
    # main data: x_value, y_value, count
    x_centers = (xedges[:-1] + xedges[1:]) / 2
    y_centers = (yedges[:-1] + yedges[1:]) / 2
    main_data = []
    for i, xv in enumerate(x_centers):
        for j, yv in enumerate(y_centers):
            main_data.append([xv, yv, int(H[i, j])])
    save_data_txt(out_txt, append_dict, main_data, [
                  f'{x_label.lower()}_value', f'{y_label.lower()}_value', 'count'], empty_note)

# 绘制直方图


def plot_histogram(mids, idx, label, out_prefix, out_dir, bins=60, break_locations=None, xlim=None):
    out_png = os.path.join(out_dir, f'{out_prefix}_dist.png')
    out_txt = os.path.join(out_dir, f'{out_prefix}_dist.txt')
    if len(mids) == 0:
        hist = np.zeros(bins)
        edges = np.linspace(0, 1, bins+1)
        empty_note = 'No breakpoints found.'
    else:
        data = mids[:, idx]
        edges = np.linspace(data.min(), data.max(), bins+1)
        hist, _ = np.histogram(data, bins=edges)
        empty_note = None
    centers = (edges[:-1] + edges[1:]) / 2
    plt.figure(figsize=(6, 4))
    plt.bar(centers, hist, width=centers[1] -
            centers[0], color='blue', alpha=0.7)
    plt.xlabel(label)
    plt.ylabel('Count')
    plt.title(f'Breakpoints {label} Distribution')
    if break_locations is not None and len(break_locations) > 0:
        for x in break_locations:
            plt.axvline(x=x, color='r', linestyle='--', alpha=0.7)
        plt.legend(['Break Location'])
    if xlim is not None:
        plt.xlim(xlim)
    plt.tight_layout()
    plt.savefig(out_png)
    plt.close()
    # 保存数据
    append_dict = {
        f'mean_{label.lower()}': np.mean(mids[:, idx]) if len(mids) else 0,
        f'std_{label.lower()}': np.std(mids[:, idx]) if len(mids) else 0,
    }
    main_data = [[centers[i], int(hist[i])] for i in range(len(hist))]
    save_data_txt(out_txt, append_dict, main_data, [
                  f'{label.lower()}_value', 'count'], empty_note)

# 主函数


def get_xyz_path_from_task_dir(task_dir):
    # 解析 bondingprob 数值
    parent = os.path.basename(os.path.dirname(task_dir))
    if parent.startswith('string=bondingprob'):
        prob = parent.split('=')[1]
        xyz_name = f'out_deform_xyz_bondingprob{prob}.xyz'
        xyz_path = os.path.join(task_dir, xyz_name)
        if os.path.exists(xyz_path):
            return xyz_path
    # 兜底：遍历找 .xyz 文件
    for f in os.listdir(task_dir):
        if f.endswith('.xyz'):
            return os.path.join(task_dir, f)
    return None


def analyze_breakpoint_distribution(task_dir, sub_dir='analysis/breakpoint_distribution_output'):
    """
    输入: N*任务目录路径
    输出: 末态断裂点分布的热力图和直方图,所有图片和数据txt保存在task_dir/sub_dir下
    """
    out_dir = os.path.join(task_dir, sub_dir)
    if not os.path.exists(out_dir):
        os.makedirs(out_dir)
    graph_series_dir = os.path.join(task_dir, 'graph_series')
    # 读取初始帧和末态帧
    files = [f for f in os.listdir(
        graph_series_dir) if f.endswith('-graph.pkl')]
    if not files:
        raise FileNotFoundError('No *-graph.pkl found in graph_series.')
    timesteps = [int(f.split('-')[0]) for f in files]
    min_t = min(timesteps)
    max_t = max(timesteps)
    init_graph_path = os.path.join(graph_series_dir, f'{min_t}-graph.pkl')
    last_graph_path = os.path.join(graph_series_dir, f'{max_t}-graph.pkl')
    with open(init_graph_path, 'rb') as f:
        init_graph = pickle.load(f)
    with open(last_graph_path, 'rb') as f:
        last_graph = pickle.load(f)
    pos_key = auto_detect_pos_key(init_graph)
    mids = extract_breakpoints_by_compare(
        init_graph, last_graph, pos_key)
    box_x = init_graph.graph['box'][0]
    # 获取 break_location
    xyz_path = get_xyz_path_from_task_dir(task_dir)
    break_location = None
    if xyz_path is not None:
        universe = mda.Universe(xyz_path, format='LAMMPSDUMP')
        _, _, break_location_ratio = gm.find_broken_frame(universe, bins=60)
        if break_location_ratio is not None:
            # 将比例值转换为实际坐标值
            break_location = break_location_ratio * box_x
    # 热力图 yz, xy, xz
    plot_heatmap(mids, 1, 2, 'Y', 'Z', 'breakpoints_yz', out_dir)
    plot_heatmap(mids, 0, 1, 'X', 'Y', 'breakpoints_xy', out_dir)
    plot_heatmap(mids, 0, 2, 'X', 'Z', 'breakpoints_xz', out_dir)
    # 直方图 x, y, z
    plot_histogram(mids, 0, 'X', 'breakpoints_x', out_dir, break_locations=[
                   break_location] if break_location is not None else None, xlim=(0, box_x))
    plot_histogram(mids, 1, 'Y', 'breakpoints_y', out_dir)
    plot_histogram(mids, 2, 'Z', 'breakpoints_z', out_dir)


def pbc_midpoint(p0, p1, boxsize):
    # p0, p1: array-like, boxsize: float or array
    delta = np.array(p1) - np.array(p0)
    # 最小镜像
    delta = delta - np.round(delta / boxsize) * boxsize
    mid = np.array(p0) + delta / 2
    # 保证mid在[0, boxsize)内
    mid = np.mod(mid, boxsize)
    return mid


if __name__ == '__main__':
    analyze_breakpoint_distribution(
        '/home/baixj/gel_network/1-crosslink_degree_space/string=bondingprob1.0/N4')
